import numpy as np
import matplotlib.pyplot as plt
import os

# 加载数据
data_path = './dataset/channel_data/channel_data.npz'
print(f"Loading data: {data_path}")
data = np.load(data_path)

# 获取信道特征数据
channel_features = data['channel_features']  # [n_windows, window_size, n_clusters, 5]
print(f"Channel features shape: {channel_features.shape}")

# 创建保存目录
save_dir = './dataset/channel_data/cluster_metrics_simple'
os.makedirs(save_dir, exist_ok=True)

# 特征名称和单位
feature_names = [
    'Horizontal Arrival Angle (degrees)', 
    'Vertical Arrival Angle (degrees)', 
    'Received Power (dB)', 
    'Horizontal Spread Angle (ASA)', 
    'Vertical Spread Angle (ZSA)'
]

# 只处理第一个窗口
window_idx = 0
window_data = channel_features[window_idx]  # [20, 25, 5]

# 获取第一个簇(索引0)在这个窗口的所有时间点上的数据
cluster_data = window_data[:, 0, :]  # [20, 5]

# 创建单独的图表，为每个特征创建一个单独的图像
for feature_idx in range(5):
    plt.figure(figsize=(10, 6))
    
    # 获取特征数据
    feature_data = cluster_data[:, feature_idx]
    
    # 绘制折线图
    plt.plot(feature_data, 'o-', linewidth=2, markersize=8, color='royalblue')
    plt.title(f'First Cluster: {feature_names[feature_idx]}', fontsize=14)
    plt.xlabel('Timestep Index', fontsize=12)
    plt.ylabel(feature_names[feature_idx], fontsize=12)
    plt.grid(True)
    
    # 添加数据标签
    for i, v in enumerate(feature_data):
        plt.text(i, v, f'{v:.2f}', ha='center', va='bottom', fontsize=9)
    
    # 设置轴的范围以留出空间给标签
    y_range = feature_data.max() - feature_data.min()
    plt.ylim(feature_data.min() - 0.1 * y_range, feature_data.max() + 0.2 * y_range)
    
    # 保存图像
    save_path = os.path.join(save_dir, f'first_cluster_feature_{feature_idx+1}.png')
    plt.tight_layout()
    plt.savefig(save_path, dpi=150)
    plt.close()
    
    print(f"Saved feature {feature_idx+1} ({feature_names[feature_idx]}) to {save_path}")

# 创建一个综合图，在同一图表上展示所有特征（归一化）
plt.figure(figsize=(12, 8))

# 归一化每个特征以便在同一图表中对比
normalized_data = np.zeros_like(cluster_data, dtype=float)
for feature_idx in range(5):
    feature_data = cluster_data[:, feature_idx]
    min_val = feature_data.min()
    max_val = feature_data.max()
    if max_val > min_val:
        normalized_data[:, feature_idx] = (feature_data - min_val) / (max_val - min_val)
    else:
        normalized_data[:, feature_idx] = feature_data - min_val

# 使用更清晰的颜色和标记
colors = ['royalblue', 'darkorange', 'forestgreen', 'darkviolet', 'crimson']
markers = ['o', 's', '^', 'D', 'x']

# 绘制所有特征的归一化曲线
for feature_idx in range(5):
    plt.plot(
        normalized_data[:, feature_idx], 
        marker=markers[feature_idx],
        linestyle='-', 
        linewidth=2, 
        markersize=8,
        color=colors[feature_idx],
        label=feature_names[feature_idx].split('(')[0].strip()
    )

plt.title('Normalized Metrics for First Cluster', fontsize=16)
plt.xlabel('Timestep Index', fontsize=14)
plt.ylabel('Normalized Value', fontsize=14)
plt.legend(fontsize=12)
plt.grid(True)
plt.tight_layout()

# 保存综合图
combined_save_path = os.path.join(save_dir, 'first_cluster_normalized_metrics.png')
plt.savefig(combined_save_path, dpi=150)
plt.close()

print(f"Saved combined normalized metrics to {combined_save_path}")
print("All visualizations complete.") 